#include <ros/ros.h>
#include <nav_msgs/Path.h>
#include <geometry_msgs/PoseStamped.h>
#include <std_msgs/Float32MultiArray.h>
#include <vector>
#include <cmath>

class ArtificialPotentialField {
public:
    ArtificialPotentialField(double k_attractive, double k_repulsive, double rho_0)
        : k_attractive(k_attractive), k_repulsive(k_repulsive), rho_0(rho_0) {}

    std::vector<double> attractive_force(const std::vector<double>& robot_pos, const std::vector<double>& goal_pos) const {
        return {k_attractive * (goal_pos[0] - robot_pos[0]), k_attractive * (goal_pos[1] - robot_pos[1])};
    }

    std::vector<double> repulsive_force(const std::vector<double>& robot_pos, const std::vector<double>& obstacle_pos) const {
        double distance = std::sqrt(std::pow(robot_pos[0] - obstacle_pos[0], 2) + std::pow(robot_pos[1] - obstacle_pos[1], 2));
        if (distance > rho_0) {
            return {0.0, 0.0};
        } else {
            return {k_repulsive * ((1.0 / distance - 1.0 / rho_0) * (robot_pos[0] - obstacle_pos[0]) / std::pow(distance, 2)),
                    k_repulsive * ((1.0 / distance - 1.0 / rho_0) * (robot_pos[1] - obstacle_pos[1]) / std::pow(distance, 2))};
        }
    }

    std::vector<double> total_force(const std::vector<double>& robot_pos, const std::vector<double>& goal_pos, const std::vector<std::vector<double>>& obstacles) const {
        std::vector<double> attractive_force = this->attractive_force(robot_pos, goal_pos);
        std::vector<double> repulsive_force = {0.0, 0.0};
        for (const auto& obstacle : obstacles) {
            repulsive_force[0] += this->repulsive_force(robot_pos, obstacle)[0];
            repulsive_force[1] += this->repulsive_force(robot_pos, obstacle)[1];
        }
        return {attractive_force[0] + repulsive_force[0], attractive_force[1] + repulsive_force[1]};
    }

private:
    double k_attractive;
    double k_repulsive;
    double rho_0;
};

void simulate_path(const std::vector<double>& robot_pos, const std::vector<std::vector<double>>& waypoints, const std::vector<std::vector<double>>& obstacles, const ArtificialPotentialField& apf, std_msgs::Float32MultiArray& route, std_msgs::Float32MultiArray& target_point) {
    std::vector<double> current_pos = robot_pos;
    for (const auto& goal_pos : waypoints) {
        for (int step = 0; step < 1000; ++step) {
            std::vector<double> force = apf.total_force(current_pos, goal_pos, obstacles);
            current_pos[0] += 0.05 * force[0];
            current_pos[1] += 0.05 * force[1];

            // Add the current position to the route
            route.data.push_back(current_pos[0]);
            route.data.push_back(current_pos[1]);

            if (std::sqrt(std::pow(current_pos[0] - goal_pos[0], 2) + std::pow(current_pos[1] - goal_pos[1], 2)) < 0.1) {
                // Add the target point
                target_point.data.push_back(goal_pos[0]);
                target_point.data.push_back(goal_pos[1]);
                break;
            }
        }
    }
}

int main(int argc, char** argv) {
    ros::init(argc, argv, "path_planning");
    ros::NodeHandle nh;

    // Publishers for route and target point
    ros::Publisher pub_route = nh.advertise<std_msgs::Float32MultiArray>("/mpc_route", 10);
    ros::Publisher pub_target = nh.advertise<std_msgs::Float32MultiArray>("/mpc_target_point", 10);

    ArtificialPotentialField apf(2.0, 3.0, 0.8);
    std::vector<double> robot_pos = {0.0, 0.0};
    std::vector<std::vector<double>> waypoints = {{4.0, 0.0}, {6.0, 0.0}, {10.0, 0.0}};
    std::vector<std::vector<double>> obstacles = {{5.0, 0.0}};

    std_msgs::Float32MultiArray route;
    std_msgs::Float32MultiArray target_point;

    simulate_path(robot_pos, waypoints, obstacles, apf, route, target_point);

    // Publish the route and target point
    pub_route.publish(route);
    pub_target.publish(target_point);

    ros::spin();

    return 0;
}
